import os.path
import random
from sklearn.model_selection import train_test_split
from keras.utils.np_utils import to_categorical
from keras.models import Sequential
from keras.layers.core import Dense
from keras.optimizers import SGD
import cv2
import numpy as np
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers import BatchNormalization
from keras.layers.convolutional import Conv2D
from keras.layers.convolutional import MaxPooling2D
from keras.layers.core import Activation
from keras.layers.core import Flatten
from keras.layers.core import Dropout
from keras.layers.core import Dense
from keras import backend as K
from imutils import paths
from sklearn.preprocessing import LabelEncoder


class NeuralNetwork:
    def __init__(self, datasetPath: str):
        self.data = []
        self.labels = []
        self.imagePaths = sorted(list(paths.list_images(datasetPath)))
        self.trainX = []
        self.trainY = []
        self.testX = []
        self.testY = []
        self.model = Sequential()
        chanDim = -1
        inputShape = (64, 64, 3)
        self.model.add(Conv2D(32, (3, 3), padding="same",
                              input_shape=inputShape))
        self.model.add(Activation("relu"))
        self.model.add(BatchNormalization(axis=chanDim))
        self.model.add(MaxPooling2D(pool_size=(2, 2)))
        self.model.add(Dropout(0.25))
        self.model.add(Conv2D(64, (3, 3), padding="same"))
        self.model.add(Activation("relu"))
        self.model.add(BatchNormalization(axis=chanDim))
        self.model.add(Conv2D(64, (3, 3), padding="same"))
        self.model.add(Activation("relu"))
        self.model.add(BatchNormalization(axis=chanDim))
        self.model.add(MaxPooling2D(pool_size=(2, 2)))
        self.model.add(Dropout(0.25))
        self.model.add(Conv2D(128, (3, 3), padding="same"))
        self.model.add(Activation("relu"))
        self.model.add(BatchNormalization(axis=chanDim))
        self.model.add(Conv2D(128, (3, 3), padding="same"))
        self.model.add(Activation("relu"))
        self.model.add(BatchNormalization(axis=chanDim))
        self.model.add(Conv2D(128, (3, 3), padding="same"))
        self.model.add(Activation("relu"))
        self.model.add(BatchNormalization(axis=chanDim))
        self.model.add(MaxPooling2D(pool_size=(2, 2)))
        self.model.add(Dropout(0.25))
        self.model.add(Flatten())
        self.model.add(Dense(512))
        self.model.add(Activation("relu"))
        self.model.add(BatchNormalization())
        self.model.add(Dropout(0.5))
        self.model.add(Dense(2))
        self.model.add(Activation("softmax"))

    def __randomizeRawData(self):
        random.seed(42)
        random.shuffle(self.imagePaths)

    def __generateData(self):
        for imagePath in self.imagePaths:
            print(imagePath)
            imgReaded = cv2.imread(imagePath)
            if imgReaded is None:
                continue
            image = cv2.resize(imgReaded, (64, 64))
            label = imagePath.split(os.path.sep)[-2]
            self.data.append(image),
            self.labels.append(label)
        self.data = np.array(self.data, dtype='float') / 255.0
        self.labels = np.array(self.labels)
        labelEncoder = LabelEncoder()
        (self.trainX, self.testX, self.trainY, self.testY) = train_test_split(self.data,
                                                                              labelEncoder.fit_transform(self.labels),
                                                                              test_size=0.25,
                                                                              random_state=42)
        self.testY = to_categorical(self.testY, num_classes=2)
        self.trainY = to_categorical(self.trainY, num_classes=2)

    def traintModel(self):
        self.__randomizeRawData()
        self.__generateData()
        INIT_LR = 0.01
        EPOCHS = 75
        BS = 32
        aug = ImageDataGenerator(rotation_range=30, width_shift_range=0.1,
                                 height_shift_range=0.1, shear_range=0.2, zoom_range=0.2,
                                 horizontal_flip=True, fill_mode="nearest")
        opt = SGD(lr=INIT_LR)
        self.model.compile(loss='binary_crossentropy', optimizer=opt, metrics=['accuracy'])
        self.model.fit(aug.flow(self.trainX, self.trainY, batch_size=BS),
                       validation_data=(self.testX, self.testY),
                       epochs=EPOCHS,
                       batch_size=32)

    def saveModel(self, path: str):
        self.model.save(path)
